package com.fortune.cat.limiter;

import com.fortune.cat.bean.RequestRateLimiterConfig;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Mono;

import java.lang.reflect.Field;
import java.math.BigDecimal;
import java.util.HashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;

import static org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter.*;

/**
 * @author zhaozeyang
 * @date 2019/4/26 14:42
 * @description
 */
@Slf4j
public class LocalCacheRateLimiter {

    private RequestRateLimiterConfig config;

    private final Map<String, com.google.common.util.concurrent.RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();


    public Mono<Response> isAllowed(String id) {
        com.google.common.util.concurrent.RateLimiter rateLimiter = rateLimiterMap.get(id);
        try {
            if (Objects.isNull(rateLimiter)) {
                rateLimiter = getRateLimiter(id);
            }
            boolean isAllowed = rateLimiter.tryAcquire(1);
            return Mono.just(new Response(isAllowed, getHeaders()));
        } catch (Exception e) {
            log.info("LocalCacheRateLimiter | error ", e);
        }
        return Mono.just(new Response(true, getHeaders()));
    }

    public Map<String, String> getHeaders() {
        Map<String, String> headers = new HashMap<>(4);
        headers.put(REMAINING_HEADER, "local cache can't count");
        headers.put(REPLENISH_RATE_HEADER,
                String.valueOf(config.getReplenishRate()));
        headers.put(BURST_CAPACITY_HEADER,
                String.valueOf(config.getBurstCapacity()));
        return headers;
    }

    private com.google.common.util.concurrent.RateLimiter getRateLimiter(String id) throws NoSuchFieldException, IllegalAccessException {
        //此处采用懒加载的方式创建每个限流规则的RateLimiter，
        //此处由于针对的key并不一定固定，如按ip限流的话，没有办法做到预加载
        synchronized (id) {
            com.google.common.util.concurrent.RateLimiter rateLimiter = rateLimiterMap.get(id);
            if (Objects.isNull(rateLimiter)) {
                Integer replenishRate = config.getReplenishRate();
                Integer burstCapacity = config.getBurstCapacity();
                rateLimiter = com.google.common.util.concurrent.RateLimiter.create(replenishRate);
                Double maxBurstSeconds = new BigDecimal((float) burstCapacity / replenishRate).setScale(2, BigDecimal.ROUND_HALF_UP).doubleValue();
                Field field = rateLimiter.getClass().getDeclaredField("maxBurstSeconds");
                field.setAccessible(true);
                field.setDouble(rateLimiter, maxBurstSeconds);
                rateLimiterMap.put(id, rateLimiter);
                return rateLimiter;
            }
        }
        return rateLimiterMap.get(id);
    }

    public void setRequestRateLimiterConfig(RequestRateLimiterConfig requestRateLimiterConfig) {
        this.config = requestRateLimiterConfig;
    }
}
